from model import ExamResultCollection, ExamCollection
from collections import OrderedDict
from argparse import ArgumentParser, Namespace
from typing import Dict, List
from os import path, listdir


# 有的题型在不同的考卷中写法不一，我们通过这个映射将题目归类整齐
alias = {
    '数学运算': '数量关系',
    '言语理解与表达': '言语理解',
    '类比推理': '判断推理',
    '逻辑判断': '判断推理',
    '图形推理': '判断推理',
    '科学推理': '判断推理',
}


def parse_args() -> Namespace:
    parser = ArgumentParser()
    parser.add_argument("exams", type=str)
    parser.add_argument("result", type=str)
    parser.add_argument("--output-csv", type=str)
    return parser.parse_args()


def collect_result(exams: ExamCollection, exam_results: ExamResultCollection) -> Dict[str, List[int]]:
    stats = OrderedDict()
    ncorr = 0
    ntotal = 0
    nreject = 0
    for exam in exams.exams:
        for question in exam.questions:
            if 'A、' not in question.content:
                # 有几个判断题没有给选项列表，我们这里直接跳过了
                continue

            ntotal += 1
            ty = question.type
            if ty in alias:
                ty = alias[ty]

            if ty not in stats:
                # 4元组：总问题数，答对数，答出数，总token消耗
                stats[ty] = [0, 0, 0, 0]

            stats[ty][0] += 1
            answer = exam_results.results[exam.id].answers.get(question.id)
            if answer is not None:
                stats[ty][3] += answer.token_cost
                if answer.correct:
                    stats[ty][1] += 1
                    ncorr += 1
            else:
                stats[ty][2] += 1

    print(f"|         | 正确数 | 总问题数 | 拒绝回答数 | 正确率 | 每问题平均token消耗 |")
    print(f"|---------|-------|--------|-----------|------|------------------|")
    for ty, (n1total, n1corr, n1reject, n1tokens) in stats.items():
        ncorr_rate = n1corr / n1total * 100
        ntoken_per = n1tokens // n1total
        nreject += n1reject
        print(f"| {ty} | {n1corr:4d} | {n1total:4d} | {n1reject:4d} | {ncorr_rate:3.1f}% | {ntoken_per:5d} |")

    print(f"| 总计     | {ncorr:4d} | {ntotal:4d} | {nreject:4d} | {ncorr / ntotal * 100:3.1f}% | |")
    stats['总计'] = [ntotal, ncorr, nreject, 0]
    return stats


def main(args: Namespace) -> None:
    with open(args.exams, 'r') as f:
        exams = ExamCollection.model_validate_json(f.read())

    result_paths = []
    if path.isfile(args.result):
        result_paths.append(args.result)
    elif path.isdir(args.result):
        for name in listdir(args.result):
            if name.endswith(".json"):
                result_paths.append(path.join(args.result, name))

    results: Dict[str, Dict[str, List[int]]] = dict()
    for result_path in result_paths:
        with open(result_path, 'r') as f:
            exam_results = ExamResultCollection.model_validate_json(f.read())

        print(f"========== {result_path:40s} ==========")
        results[result_path] = collect_result(exams, exam_results)

    if args.output_csv:
        with open(args.output_csv, 'w+', encoding='utf8') as fcsv:
            fcsv.write("模型名,题目类别,总题目数量,做对题数,拒绝题数,总token消耗\n")
            for model_path, result in results.items():
                model = model_path.split("/")[-1].replace(".json", "").lower()
                for ty, stat in result.items():
                    fcsv.write(f"{model},{ty},{stat[0]},{stat[1]},{stat[2]},{stat[3]}\n")


if __name__ == "__main__":
    args = parse_args()
    main(args)